#include "opencl_wrapper.h"
#include "opencl_helper.h"
#include "winograd_4x4_to_36.h"
#include "winograd_36_to_4x4.h"
#include "conv_1x1.h"
#include "utils.h"

#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <string>
#include <vector>
#include <fstream>
// #include <iterator>

// #define WINOGRAD_4x4_TO_36
// #define CONV_1x1
#define WINOGRAD_36_TO_4x4

int main(int argc, char const *argv[])
{
    cl_int status = 0;
    printf("------------------------------ main() start ------------------------------\n");
    mslite::loadOpenCL();
    cl_device_id device;
    cl_context context;
    cl_command_queue commandQueue;
    status = mslite::createCLEnv(1, device, context, commandQueue);

    /*********************************** WINOGRAD_4x4_TO_36 ***********************************/
    #ifdef WINOGRAD_4x4_TO_36
    std::vector<int> ishape{1, 16, 256, 100};
    std::vector<mslite::Operation *> ops = {
        new mslite::Winograd4x4To36Naive(context, device, ishape),
        new mslite::Winograd4x4To36Tflite(context, device, ishape)
    };
    #endif

    /*********************************** CONV_1x1 ***********************************/
    #ifdef CONV_1x1
    std::vector<int> ishape{1, 36, 256, 100};
    std::vector<int> wshape{96, 6, 6, 100};
    std::vector<mslite::Operation *> ops = {
        new mslite::Conv1x1Naive(context, device, ishape, wshape),
        new mslite::Conv1x1Tflite(context, device, ishape, wshape)
    };
    #endif

    /*********************************** WINOGRAD_36_TO_4x4 ***********************************/
    #ifdef WINOGRAD_36_TO_4x4
    std::vector<int> ishape{1, 36, 256, 96};
    std::vector<int> oshape{1, 16, 256, 96};
    std::vector<mslite::Operation *> ops = {
        new mslite::Winograd36To4x4Naive(context, device, ishape, oshape),
        new mslite::Winograd36To4x4Tflite(context, device, ishape, oshape)
    };
    #endif

    std::vector<std::vector<int>> ishapes = ops[0]->getInputShapes();
    std::vector<std::vector<float>> inputs(ishapes.size());
    for (int i=0; i<ishapes.size(); i++)
        mslite::genData(inputs[i], ishapes[i]);
    for (int i=0; i<ops.size(); i++) {
        ops[i]->setInput(inputs, commandQueue);
        ops[i]->bind();
        ops[i]->setParams();
    }

    // const std::vector<std::vector<size_t>> localCandidates = {{8, 6, 4}, {8, 6, 2}, {4, 6, 2}, {2, 6, 2}, 
    //                                                           {2, 6, 1}, {1, 6, 1}, {1, 3, 1}, {1, 1, 1}};

    // const std::vector<std::vector<size_t>> localCandidates = {{16, 6, 4}, {8, 6, 4}, {8, 6, 2}, {4, 6, 2}, 
    //                                                           {4, 3, 2}, {2, 6, 2}, {2, 3, 2}, {2, 3, 1},
    //                                                           {1, 3, 1}, {1, 1, 1}};

    const std::vector<std::vector<size_t>> localCandidates = {{32, 4, 2}, {16, 4, 2}, {16, 4, 1},
                                                              {8, 4, 1},  {4, 4, 1},  {2, 4, 1},
                                                              {1, 4, 1},  {1, 2, 1},  {1, 1, 1}};

    for (auto candi: localCandidates) {
        int round = 20;
        std::vector<double> times(round, 0);
        for (int i=0; i<ops.size(); i++) {
            std::cout << ops[i]->getKernelName() << "\t";
            ops[i]->setParams(candi);

            for (int j=0; j<round; j++) {
                cl_event profEvent;
                ops[i]->run(commandQueue, profEvent);
                times[j] = mslite::profileTime(profEvent, false);
            }
            std::cout << "\t\tAvg Time: " << mslite::getAvg(times, 5) << "ms" << std::endl;
        }
    }

    std::vector<std::vector<std::vector<float>>> outputs(ops.size());
    for (int i=0; i<ops.size(); i++)
        ops[i]->getOutut(commandQueue, outputs[i]);
    
    double maxDiff = 0;
    for (int i=0; i<outputs[0][0].size(); i++) {
        // if (i < 10) {
        //     std::cout << outputs[0][0][i] << "\t" << outputs[1][0][i] << std::endl;
        // }
        maxDiff = std::max(maxDiff, std::fabs(static_cast<double>(outputs[0][0][i]-outputs[1][0][i])));
    }
    std::cout << maxDiff << std::endl;

    // 资源回收
    status = mslite::releaseEnv(context, commandQueue);
    printf("------------------------------- main() over -------------------------------\n");
    return 0;
}